﻿using System;
using System.Collections.Generic;
using System.Text;
using System.IO;
using System.Text.RegularExpressions;
using Microsoft.Win32;
using System.Collections.ObjectModel;

namespace System.Peppers.Win32
{
    /// <summary>
    /// Class for loading/saving entries from a computer's host file
    /// </summary>
    public sealed class HostFile
    {
        #region Static

        private static readonly string HOSTS_FILE_PATH = null,
            LMHOSTS_FILE_PATH = null;
        private static readonly Regex _regex = new Regex(@"\s+", RegexOptions.Compiled);
        private const string HKEY_HOST_KEY = @"HKEY_LOCAL_MACHINE\SYSTEM\CurrentControlSet\Services\Tcpip\Parameters",
            HKEY_HOST_VALUE = "DataBasePath",
            HOSTS_FILE = "hosts",
            LMHOSTS_FILE = "lmhosts",
            RELATIVE_DIR = @"drivers\etc\",
            HOSTS_PATH_RELATIVE = RELATIVE_DIR + HOSTS_FILE,
            LMHOSTS_PATH_RELATIVE = RELATIVE_DIR + LMHOSTS_FILE,
            COMMENT = "#";

        static HostFile()
        {
            try
            {
                string path = Registry.GetValue(HKEY_HOST_KEY, HKEY_HOST_VALUE, null) as string;
                if (string.IsNullOrEmpty(path))
                {
                    path = Environment.GetFolderPath(Environment.SpecialFolder.System);
                }
                else
                {
                    path = Environment.ExpandEnvironmentVariables(path);
                }
                HOSTS_FILE_PATH = Path.Combine(path, HOSTS_PATH_RELATIVE);
                LMHOSTS_FILE_PATH = Path.Combine(path, LMHOSTS_PATH_RELATIVE);
            }
            //Just in case there is a weird permission error
            catch
            {
                string system = Environment.GetFolderPath(Environment.SpecialFolder.System);
                HOSTS_FILE_PATH = Path.Combine(system, HOSTS_PATH_RELATIVE);
                LMHOSTS_FILE_PATH = Path.Combine(system, LMHOSTS_PATH_RELATIVE);
            }
        }

        #endregion

        private List<HostFileEntry> _entries = new List<HostFileEntry>();
        private Dictionary<string, HostFileEntry> _entriesByIP = new Dictionary<string, HostFileEntry>(),
            _entriesByHostname = new Dictionary<string, HostFileEntry>();

        /// <summary>
        /// Constructor, defaults to HostFilePath.Hosts
        /// </summary>
        public HostFile()
            : this(HostFilePath.Hosts)
        { }

        /// <summary>
        /// Constructor, takes path of file to load
        /// </summary>
        /// <param name="path">Path of host file to load</param>
        public HostFile(HostFilePath path)
        {
            Entries = new ReadOnlyCollection<HostFileEntry>(_entries);
            EntriesByIP = _entriesByIP;
            EntriesByHostname = _entriesByHostname;

            switch (path)
            {
                case HostFilePath.LmHosts:
                    FilePath = LMHOSTS_FILE_PATH;
                    break;
                case HostFilePath.Hosts:
                default:
                    FilePath = HOSTS_FILE_PATH;
                    break;
            }
        }

        /// <summary>
        /// Full path of specified file
        /// </summary>
        public string FilePath
        {
            get;
            private set;
        }

        /// <summary>
        /// List of all entries, even comments
        /// </summary>
        public ReadOnlyCollection<HostFileEntry> Entries
        {
            get;
            private set;
        }

        /// <summary>
        /// Dictionary of entries by their IP
        /// </summary>
        public IDictionary<string, HostFileEntry> EntriesByIP
        {
            get;
            private set;
        }

        /// <summary>
        /// Dictionary of entries by their hostname
        /// </summary>
        public IDictionary<string, HostFileEntry> EntriesByHostname
        {
            get;
            private set;
        }

        /// <summary>
        /// Loads the host file, make sure to wrap in try-catch
        /// </summary>
        public void Load()
        {
            _entries.Clear();
            _entriesByIP.Clear();
            _entriesByHostname.Clear();

            if (!File.Exists(FilePath))
            {
                //Don't try to load anything
                return;
            }

            using (StreamReader reader = new StreamReader(File.OpenRead(FilePath)))
            {
                do
                {
                    HostFileEntry entry = new HostFileEntry { RawText = reader.ReadLine() };
                    if (!entry.RawText.StartsWith(COMMENT))
                    {
                        string[] split = _regex.Split(entry.RawText, 2);
                        if (split.Length > 1)
                        {
                            entry.IsEntry = true;
                            entry.IP = split[0];
                            entry.Hostname = split[1];
                        }
                    }

                    if (entry.IsEntry)
                    {
                        if (!_entriesByHostname.ContainsKey(entry.Hostname))
                        {
                            if (!_entriesByIP.ContainsKey(entry.IP))
                            {
                                _entriesByIP[entry.IP] = entry;
                            }

                            _entriesByHostname[entry.Hostname] = entry;
                            _entries.Add(entry);
                        }
                    }
                    else
                    {
                        _entries.Add(entry);
                    }
                }
                while (!reader.EndOfStream);
            }
        }

        /// <summary>
        /// Saves the hosts file, also wrap in try-catch must have admin priveleges
        /// </summary>
        public void Save()
        {
            //Call create so it will create the file if it's not there
            using (StreamWriter writer = new StreamWriter(File.Create(FilePath)))
            {
                foreach (HostFileEntry entry in _entries)
                {
                    writer.WriteLine(entry.RawText);
                }
            }
        }

        /// <summary>
        /// Adds an entry, overwriting existing members for identical ip/hostname
        /// </summary>
        /// <param name="ip">IP of the entry to add</param>
        /// <param name="hostname">Hostname of the entry to add</param>
        public void Add(string ip, string hostname)
        {
            //Remove old entry if it exists
            HostFileEntry entry = null;
            if (_entriesByIP.TryGetValue(ip, out entry))
            {
                _entries.Remove(entry);
                _entriesByIP.Remove(ip);
            }
            if (_entriesByHostname.TryGetValue(hostname, out entry))
            {
                _entries.Remove(entry);
                _entriesByHostname.Remove(hostname);
            }

            entry = new HostFileEntry
                {
                    IsEntry = true,
                    IP = ip,
                    Hostname = hostname,
                    RawText = ip + "\t" + hostname,
                };

            _entries.Add(entry);
            _entriesByIP[entry.IP] = entry;
            _entriesByHostname[entry.Hostname] = entry;
        }

        /// <summary>
        /// Returns true if the specified entry is already present
        /// </summary>
        /// <param name="ip">IP of the entry</param>
        /// <param name="hostname">Hostname of the entry</param>
        /// <returns>True if found</returns>
        public bool Contains(string ip, string hostname)
        {
            HostFileEntry entry = null;
            return _entriesByIP.TryGetValue(ip, out entry) && entry.Hostname == hostname;
        }

        /// <summary>
        /// Removes an entry
        /// </summary>
        /// <param name="ip">IP of the entry</param>
        /// <param name="hostname">Hostname of the entry</param>
        public void Remove(string ip, string hostname)
        {
            HostFileEntry entry = null;
            if (_entriesByIP.TryGetValue(ip, out entry))
            {
                _entries.Remove(entry);
                _entriesByIP.Remove(ip);
            }
            if (_entriesByHostname.TryGetValue(hostname, out entry))
            {
                _entries.Remove(entry);
                _entriesByHostname.Remove(hostname);
            }
        }
    }
}